from tensorflow import keras
import matplotlib.pyplot as plt
import numpy as np
from tensorflow_core.python.keras import layers

# 1.加载cifar10数据集，并查看数据集的形状,对数据进行归一化处理（6分）
(x_train, y_train), (x_test, y_test) = keras.datasets.cifar10.load_data()
x_train = x_train.astype("float32") / 255.0
x_test = x_test.astype("float32") / 255.0

print("y_train[0]: ", y_train[0])
# 将整型的类别标签转为onehot编码。y_train为int数组
y_train = keras.utils.to_categorical(y_train, 10)
y_test = keras.utils.to_categorical(y_test, 10)
print("one hot y_train[0]: ", y_train[0])
print("x_train.shape: ", x_train.shape)
print("y_train.shape: ", y_train.shape)
print("x_test.shape: ", x_test.shape)
print("y_test.shape: ", y_test.shape)

# 可视化训练集前16张图片，并在图像下方显示类别
label_dict = {0: "airplane", 1: "automobile", 2: 'bird', 3: 'cat', 4: 'deer', 5: 'dog', 6: 'frog',
              7: 'horse', 8: 'ship', 9: 'truck'}
plt.figure(figsize=(10, 10))
for i in range(16):
    plt.subplot(4, 4, i + 1)
    plt.imshow(x_train[i])
    plt.xlabel(label_dict[np.argmax(y_train[i])])
plt.tight_layout()  # 优化布局
plt.show()

# 2.构建resnet模型进行训练，模型准确率达到75%左右（16分）
inputs = keras.Input(shape=(32, 32, 3))
x = layers.Conv2D(32, 3, activation="relu")(inputs)
x = layers.Conv2D(64, 3, activation="relu")(x)
block_1_output = layers.MaxPooling2D(3)(x)

x = layers.Conv2D(64, 3, activation="relu", padding="same")(block_1_output)
x = layers.Conv2D(64, 3, activation="relu", padding="same")(x)
block_2_output = layers.add([x, block_1_output])

x = layers.Conv2D(64, 3, activation="relu", padding="same")(block_2_output)
x = layers.Conv2D(64, 3, activation="relu", padding="same")(x)
block_3_output = layers.add([x, block_2_output])

x = layers.Conv2D(64, 3, activation="relu")(block_3_output)
x = layers.GlobalAveragePooling2D()(x)
x = layers.Dense(256, activation="relu")(x)
outputs = layers.Dense(10)(x)

model = keras.Model(inputs, outputs, name="cifar10_resnet")

# 3.查看模型结构（2分）
model.summary()

# 4.保存模型图片,命名为resnet-cifar.png（2分）
keras.utils.plot_model(model, "resnet-cifar.png", show_shapes=True)

# 编译模型
model.compile(
    optimizer=keras.optimizers.Adam(),
    loss=keras.losses.CategoricalCrossentropy(from_logits=True),
    metrics=["acc"]
)
# 训练模型
history = model.fit(x_train, y_train, batch_size=64, epochs=10, validation_split=0.2)

# 评估模型
test_loss, test_acc = model.evaluate(x_test, y_test, verbose=1)
print("\nTest accuracy：", test_acc)

# 5.保存模型，模型名字为resnet_model（2分）
model.save("resnet_model")

# 6.绘制模型的准确率以及损失曲线（16分）
acc = history.history["acc"]
val_acc = history.history["val_acc"]

loss = history.history["loss"]
val_loss = history.history["val_loss"]

plt.figure(figsize=(8, 8))
plt.subplot(211)
plt.plot(acc, label="Training Accuracy")
plt.plot(val_acc, label="Valitation Accuracy")
plt.legend(loc="lower right")  # 图例显示的位置
plt.ylabel("Accuracy")
plt.ylim([min(plt.ylim()), 1])
plt.title("Training and Validation Accuracy")

plt.subplot(212)
plt.plot(loss, label="Training Loss")
plt.plot(val_loss, label="Validation Loss")
plt.legend(loc="upper right")
plt.ylabel("Cross Entropy")
plt.ylim([min(plt.ylim()), 2.0])
plt.title("Training and Validation Loss")
plt.xlabel("epoch")
plt.show()
